#!/usr/bin/env python
# -*- coding: utf-8 -*-

#PBS -N job_name
#PBS -o job.log
#PBS -e job.err
#PBS -q gold5120
#PBS -l nodes=1:ppn=8
#PBS -l walltime=22:00:00

import os
from ase import Atoms, Atom
from ase.io import write
from ase.optimize import MDMin
from ase.neb import NEB
from ase.neb import NEBTools
from ase.calculators.abacus.abacus_out import Abacus

os.environ['ASE_ABACUS_COMMAND'] = "mpirun -np 8 /home/shenzx/software/abacus/abacus_v2.0/bin/ABACUS.mpi.2.0>>PREFIX.log"
WorkPath = "/home/shenzx/project/python_20190718/v20200729/ase-abacus/example/neb/" 


# Create a structure
os.chdir(WorkPath)
StruCell = [4.0614, 2.8718, 2.8718, 90, 90, 90]
StruPositions = [[0.0, 0.0, 0.0],
                 [2.0307, 1.4359, -1.4359]]
initial = Atoms('Al2',
                positions=StruPositions,
                cell=StruCell,
                pbc=(1, 1, 0))

initial.append(Atom('C', (2.0307, 1.4359, 4.3077)))
initial.center(vacuum=4.0, axis=2)

final = initial.copy()
final.positions[-1][1] += 2.8718
final.positions[-1][0] += 4.0614

# Construct a list of images:
images = [initial]
for i in range(3):
    images.append(initial.copy())
images.append(final)
print('Create images successfully')

for image in images:
    # Let all images use an abacus calculator:
    image.set_calculator(Abacus(
          atoms=image,
          pseudo_dir="/home/shenzx/software/abacus/SG15_ONCV_PBE_1.0",
          potential_name="PotSG15",
          basis_dir="/home/shenzx/software/abacus/Orb_DZP_E100_Standard_v1.0",
          basis_name="SG15act",
          calculation='scf',
          ntype=2,
          nbands=20,
          ecutwfc=50,
          dr2="1.0e-6",
          niter=100,
          force=1,
          smearing='gaussian',
          sigma=0.02,
          mixing_type='pulay-kerker',
          ks_solver='genelpa',
          mixing_beta=0.4,
          basis_type='lcao',
          atom_file='STRU',
          gamma_only=0,
          knumber=0,
          kmode='Gamma',
          kpts=[2, 2, 2, 0, 0, 0]
          ))

# Create a Nudged Elastic Band:
neb = NEB(images)
# Make a starting guess for the minimum energy path (a straight line
# from the initial to the final state):
neb.interpolate()
# Relax the NEB path:
minimizer = MDMin(neb)
print('Relax the NEB path, please wait!!!')
minimizer.run(fmax=0.5)
# Write the path to a trajectory:
write('neb.traj', images)
print('Write the path to a trajectory successfully')
nebtools = NEBTools(images)
# Get the calculated barrier and the energy change of the reaction.
Ef, dE = nebtools.get_barrier()
print("barrier        Energy of reaction")
print(Ef, '   ', dE)
# Get the actual maximum force at this point in the simulation.
max_force = nebtools.get_fmax()
print('max_force:   ', max_force)
